

#include"undistort.h"

#include <ros/ros.h>
#include "rosbag/bag.h"
#include "rosbag/view.h"
#include <nav_msgs/Odometry.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/PointCloud2.h>

#include "pcl_conversions/pcl_conversions.h"
#include <pcl/io/ply_io.h>
#include <pcl/registration/icp.h>

#include <pclomp/ndt_omp.h>
#include <pclomp/gicp_omp.h>
#include <pcl/registration/icp.h>
#include <pcl/registration/gicp.h>
#include <pcl/registration/ndt.h>

#include"utils.h"
#include"cali.h"
#include"ndtSlam.h"

#define  MIN_ANGLE_DIS 1
#define  MAX_ANGLE_DIS 5


Eigen::Quaterniond solve(const vector<pair<Eigen::Quaterniond, Eigen::Quaterniond>> &corres)
{
    if (corres.size() == 0)
    {
        cout << "no constraint found !!!" << endl;
        return move(Eigen::Quaterniond().Identity());
    }

    cout << "constraints size " << corres.size() << endl;

    // transform quaternion to skew symmetric matrix
    auto toSkewSymmetric = [](const Eigen::Vector3d &q) -> Eigen::Matrix3d {
        Eigen::Matrix3d mat = Eigen::Matrix3d::Zero();
        mat(0, 1) = -q.z();
        mat(0, 2) = q.y();
        mat(1, 0) = q.z();
        mat(1, 2) = -q.x();
        mat(2, 0) = -q.y();
        mat(2, 1) = q.x();

        return move(mat);
    };

    // create homogeneous linear equations
    Eigen::MatrixXd A(corres.size() * 4, 4);
    for (int i = 0; i < corres.size(); i++)
    {
        // get relative transform
        const auto &q_l2_l1 = corres[i].first;
        const auto &q_b2_b1 = corres[i].second;

        // get left product matrix
        Eigen::Vector3d q_b2_b1_vec = q_b2_b1.vec();
        Eigen::Matrix4d left_Q_b2_b1 = Eigen::Matrix4d::Zero();
        left_Q_b2_b1.block<1, 3>(0, 1) = -q_b2_b1_vec.transpose();
        left_Q_b2_b1.block<3, 1>(1, 0) = q_b2_b1_vec;
        left_Q_b2_b1.block<3, 3>(1, 1) = toSkewSymmetric(q_b2_b1_vec);
        left_Q_b2_b1 += q_b2_b1.w() * Eigen::Matrix4d::Identity();

        // get right product matrix
        Eigen::Vector3d q_l2_l1_vec = q_l2_l1.vec();
        Eigen::Matrix4d right_Q_l2_l1 = Eigen::Matrix4d::Zero();
        right_Q_l2_l1.block<1, 3>(0, 1) = -q_l2_l1_vec.transpose();
        right_Q_l2_l1.block<3, 1>(1, 0) = q_l2_l1_vec;
        right_Q_l2_l1.block<3, 3>(1, 1) = -toSkewSymmetric(q_l2_l1_vec);
        right_Q_l2_l1 += q_l2_l1.w() * Eigen::Matrix4d::Identity();

        // add loss function
        double angle_distance = 180.0 / M_PI * q_b2_b1.angularDistance(q_b2_b1);
        double huber = angle_distance > 5.0 ? 5.0 / angle_distance : 1.0;

        A.block<4, 4>(i * 4, 0) = huber * (left_Q_b2_b1 - right_Q_l2_l1);
    }

    // solve homogeneous linear equations by svd method
    Eigen::JacobiSVD<Eigen::MatrixXd> svd(A, Eigen::ComputeFullU | Eigen::ComputeFullV);
    Eigen::Matrix<double, 4, 1> x = svd.matrixV().col(3);
    Eigen::Quaterniond q_l_b(x(0), x(1), x(2), x(3)); //lidar to body(imu), 激光到惯导的标定矩阵

    return move(q_l_b);
}



CImuLidarCali::CImuLidarCali()
{
    mRotFromImu2Lidar.setIdentity();
}

CImuLidarCali::~CImuLidarCali()
{

}

int CImuLidarCali::InitCali(string bag_filename, string lidarTopic, string imuTopic)
{
    pcl::PLYWriter ply_writer;

    //string bag_filename = "/home/xdh/data/bag/lio-sam/identity/rotation.bag";
    std::cout<<bag_filename<<std::endl;

    pclomp::NormalDistributionsTransform<PT_TYPE, PT_TYPE> ompNdt;
    int num_threads = omp_get_max_threads();
    printf("max threads: %d \n", num_threads);
    ompNdt.setResolution(1.0);
    ompNdt.setNumThreads(num_threads-1);
    ompNdt.setNeighborhoodSearchMethod(pclomp::DIRECT7);
    // ompNdt.setNeighborhoodSearchMethod(pclomp::KDTREE);
    rosbag::Bag bag;
    bag.open(bag_filename, rosbag::bagmode::Read);
    rosbag::View view(bag);
    const ::ros::Time begin_time = view.getBeginTime();
    const double duration_in_seconds =
        (view.getEndTime() - begin_time).toSec();
    cout << duration_in_seconds << endl;

    CLidarUnDistort undistort;

    vector<LidarFrame> vecLidarFrame;
    vector<ImuData>    vecImuData;
    std::deque<sensor_msgs::PointCloud2> vecCloudMsg;
    for (const rosbag::MessageInstance &message : view)
    {
        std::string stopic = message.getTopic();

        if (message.isType<sensor_msgs::PointCloud2>() && stopic==lidarTopic )
        {
            const sensor_msgs::PointCloud2::ConstPtr &msg =
                message.instantiate<sensor_msgs::PointCloud2>();

            double initStamp = msg->header.stamp.toSec();

            pcl::PointCloud<PointXYZIRT> pcl_point_cloud;
            pcl::fromROSMsg(*msg, pcl_point_cloud);
            //ply_writer.write("pc-raw.ply", pcl_point_cloud);

            pcl::PointCloud<PT_TYPE>  pointCloud;
            for (int i = 0; i < pcl_point_cloud.size(); i++)
            {
                PT_TYPE pt;
                pt.x = pcl_point_cloud.points[i].x;
                pt.y = pcl_point_cloud.points[i].y;
                pt.z = pcl_point_cloud.points[i].z;
                pointCloud.push_back(pt);
            }
            //voxel filter
            //VoxelFilter(pointCloud, 0.1);

            LidarFrame lf;
            lf.stamp = initStamp;
            lf.cloud = pointCloud;
            vecLidarFrame.push_back(lf);

            if (vecLidarFrame.size() > 100)
                break;
        }

        if( message.isType<sensor_msgs::Imu>() && stopic==imuTopic)
        {
             const sensor_msgs::Imu::ConstPtr &msg =
                message.instantiate<sensor_msgs::Imu>();

            ImuData data;
            data.acc = Eigen::Vector3d(msg->linear_acceleration.x,
                                       msg->linear_acceleration.y,
                                       msg->linear_acceleration.z);
            data.gyr = Eigen::Vector3d(msg->angular_velocity.x,
                                       msg->angular_velocity.y,
                                       msg->angular_velocity.z);
            data.stamp = msg->header.stamp.toSec();
            vecImuData.push_back(data);
            undistort.AddImu(data);
        }
    }

    //对点云进行匹配, 生成旋转角度的约束
    vector<pair<Eigen::Quaterniond, Eigen::Quaterniond>> corres(0);
    for(int i=0; i<vecLidarFrame.size()-1; i++)
    {
        double l1_time = vecLidarFrame[i].stamp;
        double l2_time = vecLidarFrame[i + 1].stamp;

        if(l1_time<vecImuData[0].stamp)
            continue;    

        ompNdt.setInputTarget(vecLidarFrame[i].cloud.makeShared());
        ompNdt.setInputSource(vecLidarFrame[i+1].cloud.makeShared());

        pcl::PointCloud<PT_TYPE>::Ptr aligned(new pcl::PointCloud<PT_TYPE>());
        ompNdt.align(*aligned); 
        float score = ompNdt.getFitnessScore();
        Eigen::Matrix4f mt = ompNdt.getFinalTransformation();
        Eigen::Matrix4d dmt = mt.cast<double>();

        // use icp to align
        // pcl::IterativeClosestPoint<PT_TYPE, PT_TYPE> icp;
        // icp.setInputTarget(vecLidarFrame[i].cloud.makeShared());
        // icp.setInputSource(vecLidarFrame[i+1].cloud.makeShared());

        // pcl::PointCloud<PT_TYPE>::Ptr aligned(new pcl::PointCloud<PT_TYPE>());
        // icp.align(*aligned); 
        // float score = icp.getFitnessScore();
        // Eigen::Matrix4f mt = icp.getFinalTransformation();
        // Eigen::Matrix4d dmt = mt.cast<double>();        

        // use gicp to align
        // pcl::GeneralizedIterativeClosestPoint<PT_TYPE, PT_TYPE> gicp;
        // gicp.setInputTarget(vecLidarFrame[i].cloud.makeShared());
        // gicp.setInputSource(vecLidarFrame[i+1].cloud.makeShared());

        // pcl::PointCloud<PT_TYPE>::Ptr aligned(new pcl::PointCloud<PT_TYPE>());
        // gicp.align(*aligned); 
        // float score = gicp.getFitnessScore();
        // Eigen::Matrix4f mt = gicp.getFinalTransformation();
        // Eigen::Matrix4d dmt = mt.cast<double>();     

        // use ndt to align
        // pcl::NormalDistributionsTransform<PT_TYPE, PT_TYPE> ndt;
        // ndt.setResolution(1.0);
        // ndt.setInputTarget(vecLidarFrame[i].cloud.makeShared());
        // ndt.setInputSource(vecLidarFrame[i+1].cloud.makeShared());

        // pcl::PointCloud<PT_TYPE>::Ptr aligned(new pcl::PointCloud<PT_TYPE>());
        // ndt.align(*aligned); 
        // float score = ndt.getFitnessScore();
        // Eigen::Matrix4f mt = ndt.getFinalTransformation();
        // Eigen::Matrix4d dmt = mt.cast<double>();
        
        if (ompNdt.hasConverged()) //&& score<1.0 )
        // if (ndt.hasConverged()) //&& score<1.0 )
        {
            Eigen::Quaterniond q_l2_l1 = Eigen::Quaterniond( dmt.block<3, 3>(0, 0)); //lidar

            //计算imu相对旋转
            double ax,ay,az;
            undistort.FindAngle(0, l1_time, ax, ay ,az);
            Eigen::Affine3f rotation = pcl::getTransformation(0, 0, 0, ax, ay, az);
            Eigen::Matrix3f m = (rotation.matrix()).block<3,3>(0,0);
            Eigen::Quaterniond q1 = Eigen::Quaterniond(m.cast<double>());

            undistort.FindAngle(0, l2_time, ax, ay ,az);
            rotation = pcl::getTransformation(0, 0, 0, ax, ay, az);
            m = (rotation.matrix()).block<3,3>(0,0);
            Eigen::Quaterniond q2 = Eigen::Quaterniond(m.cast<double>());
            Eigen::Quaterniond q_b2_b1 = q1.inverse()*q2; //imu
            double angle_distance = q2.angularDistance(q1) / 3.1415926 * 180;

            //计算imu相对旋转
            Eigen::Quaterniond cq1,cq2;
            undistort.InterpolateQuaternion(0, l1_time, cq1);
            undistort.InterpolateQuaternion(0, l2_time, cq2);
            Eigen::Quaterniond q_b2_b1_c = cq1.inverse()*cq2; //imu
            std::cout<<i<<"...."<<"score: "<<score<<" angle dis: "<<angle_distance<<std::endl; 
            corres.push_back(move(pair<Eigen::Quaterniond, Eigen::Quaterniond>(q_l2_l1, q_b2_b1_c)));
        }
    }

    Eigen::Quaterniond q_l_b_ = solve(corres);
    Eigen::Quaterniond bench_cali;
    bench_cali.setIdentity();
    // Eigen::Matrix3d rotation_matrix;
    // rotation_matrix<<-0.0261976,-0.999639,0.00601718,-0.999651,0.0261769,-0.00348948,0.00333071,-0.00610649,-0.999976;
    // Eigen::Quaterniond bench_cali(rotation_matrix);
    double diffAngle = (q_l_b_.inverse()).angularDistance(bench_cali);
    printf("cali diff: %lf \n", diffAngle/3.14159*180);

    tf::Matrix3x3 mat(tf::Quaternion(q_l_b_.x(), q_l_b_.y(), q_l_b_.z(), q_l_b_.w()));
    double roll,pitch,yaw;
    mat.getRPY(roll, pitch, yaw);
    double dpi = 1.0/3.1415926*180;
    Eigen::Matrix3d rot = Eigen::Matrix3d(Eigen::AngleAxisd(yaw, Eigen::Vector3d::UnitZ()) * Eigen::AngleAxisd(pitch, Eigen::Vector3d::UnitY()) * Eigen::AngleAxisd(roll, Eigen::Vector3d::UnitX()));
    cout<<"rpy:"<<roll*dpi<<" "<<pitch*dpi<<" "<<yaw*dpi<<std::endl;
    cout << "rotation:"<<" \n " << rot.inverse() << endl;
    
    mRotFromImu2Lidar = rot.inverse();

    return 0;
}

int CImuLidarCali::RefineCali(string bagfile, string lidarTopic, string imuTopic)
{
    
    pcl::PLYWriter ply_writer;
    std::cout<<bagfile<<std::endl;

    pclomp::NormalDistributionsTransform<PT_TYPE, PT_TYPE> pc_reg;
    int num_threads = omp_get_max_threads();
    printf("max threads: %d \n", num_threads);
    pc_reg.setResolution(1.0);
    pc_reg.setNumThreads(num_threads-1);
    pc_reg.setNeighborhoodSearchMethod(pclomp::DIRECT7);
    //pc_reg.setNeighborhoodSearchMethod(pclomp::KDTREE);
    // pcl::IterativeClosestPoint<PT_TYPE, PT_TYPE> pc_reg;
    // pcl::GeneralizedIterativeClosestPoint<PT_TYPE, PT_TYPE> pc_reg;
    pcl::NormalDistributionsTransform<PT_TYPE, PT_TYPE> pc_reg;
    // pc_reg.setResolution(1.0);
    
    

    rosbag::Bag bag;
    bag.open(bagfile, rosbag::bagmode::Read);
    rosbag::View view(bag);
    const ::ros::Time begin_time = view.getBeginTime();
    const double duration_in_seconds =
        (view.getEndTime() - begin_time).toSec();
    cout << duration_in_seconds << endl;

    CLidarUnDistort undistort;
    CImuIntegration imuIntegrate;

    vector<LidarFrame> vecLidarFrame;
    vector<ImuData>    vecImuData;
    std::deque<sensor_msgs::PointCloud2> vecCloudMsg;
    for (const rosbag::MessageInstance &message : view)
    {
        std::string stopic = message.getTopic();

        if (message.isType<sensor_msgs::PointCloud2>() && stopic==lidarTopic)
        {
            const sensor_msgs::PointCloud2::ConstPtr &msg =
                message.instantiate<sensor_msgs::PointCloud2>();

            vecCloudMsg.push_back(*msg);

            if (vecCloudMsg.size() > 5)
            {
                sensor_msgs::PointCloud2 message = vecCloudMsg.front();
                vecCloudMsg.pop_front();

                double initStamp = message.header.stamp.toSec();

                pcl::PointCloud<PointXYZIRT> pcl_point_cloud;
                pcl::fromROSMsg(message, pcl_point_cloud);

                //对当前点云进行运动畸变校正
                int res = undistort.Process(pcl_point_cloud, initStamp);
                //将校正成功的点云进行slam处理
                if (res >= 0)
                {
                    pcl::PointCloud<PT_TYPE> deskewedPointCloud;
                    for (int i = 0; i < pcl_point_cloud.size(); i++)
                    {
                        PT_TYPE pt;
                        pt.x = pcl_point_cloud.points[i].x;
                        pt.y = pcl_point_cloud.points[i].y;
                        pt.z = pcl_point_cloud.points[i].z;
                        deskewedPointCloud.push_back(pt);
                    }
                    LidarFrame lf;
                    lf.stamp = initStamp;
                    lf.cloud = deskewedPointCloud;
                    vecLidarFrame.push_back(lf);

                    if(vecLidarFrame.size()>100)
                        break;
                }
            }
        }

        if( message.isType<sensor_msgs::Imu>() && stopic==imuTopic)
        {
             const sensor_msgs::Imu::ConstPtr &msg =
                message.instantiate<sensor_msgs::Imu>();

            ImuData data;
            data.acc = Eigen::Vector3d(msg->linear_acceleration.x,
                                       msg->linear_acceleration.y,
                                       msg->linear_acceleration.z);
            data.gyr = Eigen::Vector3d(msg->angular_velocity.x,
                                       msg->angular_velocity.y,
                                       msg->angular_velocity.z);
            data.stamp = msg->header.stamp.toSec();

            imuIntegrate.AddImu(data);

            //根据初始标定的矩阵来对imu数据进行转换
            data.gyr = mRotFromImu2Lidar*data.gyr;
            vecImuData.push_back(data);
            undistort.AddImu(data);
        }
    }

    //对点云进行匹配, 生成旋转角度的约束
    vector<pair<Eigen::Quaterniond, Eigen::Quaterniond>> corres(0);
    for(int i=0; i<vecLidarFrame.size()-1; i++)
    {
        double l1_time = vecLidarFrame[i].stamp;
        double l2_time = vecLidarFrame[i + 1].stamp;

        if(l1_time<vecImuData[0].stamp)
            continue;    

        //根据imu的数据来预测相邻两帧点云之间的旋转矩阵
        double ax, ay, az;
        undistort.FindAngle(0, l1_time, ax, ay, az);
        Eigen::Affine3f ml1 = pcl::getTransformation(0, 0, 0, ax, ay, az);
        undistort.FindAngle(0, l2_time, ax, ay ,az);
        Eigen::Affine3f ml2 = pcl::getTransformation(0, 0, 0, ax, ay, az);
        Eigen::Matrix4f mr = (ml1.matrix()).inverse()*(ml2.matrix());
        
        //如果不需要预测，就设置为单位阵
        // mr.setIdentity();

        //利用点云匹配来计算相对旋转
        pc_reg.setInputTarget(vecLidarFrame[i].cloud.makeShared());
        pc_reg.setInputSource(vecLidarFrame[i+1].cloud.makeShared());


        //输出临时的点云结果
        pcl::PointCloud<PT_TYPE> predictPointCloud;
        pcl::PointCloud<PT_TYPE>::Ptr aligned(new pcl::PointCloud<PT_TYPE>());
        pc_reg.align(*aligned, mr);
        float score = pc_reg.getFitnessScore();
        Eigen::Matrix4f mt = pc_reg.getFinalTransformation();      
        Eigen::Matrix4d dmt = mt.cast<double>();
        if(pc_reg.hasConverged())
        {
            Eigen::Quaterniond q_l2_l1 = Eigen::Quaterniond( dmt.block<3, 3>(0, 0)); //lidar

            //计算imu相对旋转
            imuIntegrate.FindAngle(0, l1_time, ax, ay ,az);
            Eigen::Affine3f rotation = pcl::getTransformation(0, 0, 0, ax, ay, az);
            Eigen::Matrix3f m = (rotation.matrix()).block<3,3>(0,0);
            Eigen::Quaterniond q1 = Eigen::Quaterniond(m.cast<double>());

            imuIntegrate.FindAngle(0, l2_time, ax, ay ,az);
            rotation = pcl::getTransformation(0, 0, 0, ax, ay, az);
            m = (rotation.matrix()).block<3,3>(0,0);
            Eigen::Quaterniond q2 = Eigen::Quaterniond(m.cast<double>());
            Eigen::Quaterniond q_b2_b1 = q1.inverse()*q2; //imu

            double angle_distance = q2.angularDistance(q1) / 3.1415926 * 180;
            //计算imu相对旋转
            Eigen::Quaterniond cq1,cq2;
            imuIntegrate.InterpolateQuaternion(0, l1_time, cq1);
            imuIntegrate.InterpolateQuaternion(0, l2_time, cq2);
            Eigen::Quaterniond q_b2_b1_c = cq1.inverse()*cq2; //imu
            std::cout<<i<<"...."<<"score: "<<score<<" angle dis: "<<angle_distance<<std::endl; 
            corres.push_back(move(pair<Eigen::Quaterniond, Eigen::Quaterniond>(q_l2_l1, q_b2_b1_c)));
        }
    }

    Eigen::Quaterniond q_l_b_ = solve(corres);
    Eigen::Quaterniond bench_cali;
    bench_cali.setIdentity();
    // Eigen::Matrix3d  rotation_matrix;
    // rotation_matrix<<-0.0261976,-0.999639,0.00601718,-0.999651,0.0261769,-0.00348948,0.00333071,-0.00610649,-0.999976;
    // Eigen::Quaterniond bench_cali(rotation_matrix);
    double diffAngle = (q_l_b_.inverse()).angularDistance(bench_cali);
    printf("cali diff: %lf \n", diffAngle/3.14159*180);

    tf::Matrix3x3 mat(tf::Quaternion(q_l_b_.x(), q_l_b_.y(), q_l_b_.z(), q_l_b_.w()));
    double roll,pitch,yaw;
    mat.getRPY(roll, pitch, yaw);
    double dpi = 1.0/3.1415926*180;
    Eigen::Matrix3d rot = Eigen::Matrix3d(Eigen::AngleAxisd(yaw, Eigen::Vector3d::UnitZ()) * Eigen::AngleAxisd(pitch, Eigen::Vector3d::UnitY()) * Eigen::AngleAxisd(roll, Eigen::Vector3d::UnitX()));
    cout<<"rpy:"<<roll*dpi<<" "<<pitch*dpi<<" "<<yaw*dpi<<std::endl;
    cout << "rotation:"<<" \n " << rot.inverse() << endl;
    
    mRotFromImu2Lidar = rot.inverse();

    return 0;
}